{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "\n",
    "from utils import *\n",
    "\n",
    "# 设置数据集的路径和下载选项\n",
    "Dataset_dir = '../data'\n",
    "Output_dir = './results'\n",
    "train_metrics_dir = os.path.join(Output_dir, 'train_metrics.csv')\n",
    "test_metrics_dir = os.path.join(Output_dir, 'test_metrics.csv')\n",
    "#如果输出目录不存在，则创建\n",
    "if not os.path.exists(Output_dir):\n",
    "    os.makedirs(Output_dir)\n",
    "\n",
    "#定义数据预处理\n",
    "transform = transforms.Compose([transforms.ToTensor()]) \n",
    "# 读取CITAR-100数据集\n",
    "trainset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=True, download=True, transform=transform)\n",
    "testset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=False, download=True, transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_img_metrics(cifar100_dataset, Output_dir):\n",
    "    print(\"Calculating image metrics...\")\n",
    "    # 定义将张量转换为PIL图像的函数\n",
    "    to_pil = transforms.ToPILImage()\n",
    "    results_list = []  # 使用列表来收集数据\n",
    "\n",
    "    # 定义列名\n",
    "    columns = ['index', 'cifar100_label', 'SNR', 'EN', 'SF', 'SD', 'AG']\n",
    "\n",
    "    label_one_classes = [13, 20, 56, 24, 25, 42, 67, 0, 36, 47]\n",
    "    label_zero_classes = [7, 37, 4, 12, 79, 71, 55, 3, 40, 63]\n",
    "\n",
    "    # 使用tqdm来显示进度\n",
    "    for i in tqdm(range(len(cifar100_dataset))):\n",
    "        # 读取图像数据(张量形式)\n",
    "        image_ten, label = cifar100_dataset[i]\n",
    "        if label in label_one_classes:\n",
    "            new_value = 1\n",
    "        elif label in label_zero_classes:\n",
    "            new_value = 0\n",
    "        else:\n",
    "            continue\n",
    "        \n",
    "        # 将张量转为PIL图像\n",
    "        image_img = to_pil(image_ten)\n",
    "        # 将PIL图像转为灰度图像\n",
    "        gray_img = image_img.convert('L')   \n",
    "        # 将灰度图像转为numpy数组\n",
    "        img_int = np.array(gray_img).astype(np.int32)\n",
    "        img_double = np.array(gray_img).astype(np.float32)\n",
    "\n",
    "        # 计算信息熵EN\n",
    "        EN = calculate_EN(img_int)\n",
    "        # 计算空间频率SF\n",
    "        SF = calculate_SF(img_double)\n",
    "        # 计算标准差SD\n",
    "        SD = calculate_SD(img_double)\n",
    "        # 计算平均梯度\n",
    "        AG = calculate_AG(img_double)\n",
    "        # 计算信噪比SNR\n",
    "        SNR = calculate_SNR(image_ten)\n",
    "\n",
    "        # 将结果添加到列表中\n",
    "        results_list.append({\n",
    "            'index': i,\n",
    "            'cifar100_label': label,\n",
    "            'SNR': SNR,\n",
    "            'EN': EN,\n",
    "            'SF': SF,\n",
    "            'SD': SD,\n",
    "            'AG': AG\n",
    "        })\n",
    "\n",
    "    print(\"Saving file...\")\n",
    "    # 将列表转换为DataFrame，并指定列名\n",
    "    results_df = pd.DataFrame(results_list, columns=columns)\n",
    "\n",
    "    # 保存结果到csv文件中   \n",
    "    csv_file_path = Output_dir\n",
    "    results_df.to_csv(csv_file_path, index=False)\n",
    "    print(\"Done.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating image metrics...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50000/50000 [00:09<00:00, 5392.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving file...\n",
      "Done.\n",
      "Calculating image metrics...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:01<00:00, 5997.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving file...\n",
      "Done.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "generate_img_metrics(trainset,train_metrics_dir)\n",
    "generate_img_metrics(testset,test_metrics_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
